from operator import truediv
from platform import version
from typing import final
from urllib.parse import uses_fragment

import paramiko
import time
import re
import logging
from datetime import datetime

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('HUAWEI_Config.log'),
        logging.StreamHandler()
    ]

)
logger = logging.getLogger('HUAWEI_Config')

class HUAWEI_Config():
    def __init__(self, host, user, passwd, port=22):
        self.host = host
        self.user = user
        self.passwd = passwd
        self.port = port
        self.client = None
        self.channel = None
        self.version = None
        self.model = None
        self.command_map = {}
    def connect(self):
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.client.connect(
                hostname=self.host,
                port=self.port,
                username=self.user,
                password=self.passwd,
                timeout=10,
                banner_timeout=10,
                allow_agent=False,
                look_for_keys=False
            )

            self.channel = self.client.invoke_shell()
            time.sleep(1)

            self._get_device_info()
            self._detect_command_variations()

            logger.info(f"成功连接{self.host}({self.model},VRP{self.version})")
            return True
        except Exception as e:
            logger.error(f"连接失败:{str(e)}")
            return False

    def _get_device_info(self):
        self._send_command("display version")
        output = self._receive_output()

        version_match=re.search(r'VRP \(R\) software, Version(\S+)', output)
        self.version = version_match.group(1) if version_match else "Unknown"

        model_match=re.search(r'^.*?(S[0-9]+[A-Z]*)\s', output,re.MULTILINE)
        self.model = model_match.group(1) if model_match else "Unknown"
    def _detect_command_variations(self):
        self.command_map = {
            'sys':'system-view',
            'quit':'quit',
            'save':'save'
        }

        self._send_command("display current-configuration | include password policy")
        output = self._receive_output()

        if "local-aaa-user password policy" in output:
            self.command_map['password policy'] = "local-aaa-user password policy"
        elif "password policy" in output:
            self.command_map['password policy'] = "password policy"
        else:
            self.command_map['password policy'] = "lcoal-aaa-user password policy"

        self._send_command("rsa local-key-pair ?")
        output = self._receive_output()

        if "create" in output:
            self.command_map['rsa_create'] = "rsa local-key-pair create"
        elif "generate" in output:
            self.command_map['rsa-create'] = "rsa local-key-pair generate"
        else:
            self.command_map['rsa-create'] = "rsa local-key-pair create"

    def _send_command(self, command, wait = 0.5):
        if not self.channel:
            return False
        final_command = self.command_map.get(command, command)
        self.channel.send(final_command+"\n")
        time.sleep(wait)
        return True
    def _receive_output(self,timeout = 5):

        if not self.channel:
            return ""

        output = ""

        start_time = time.time()

        while time.time() - start_time < timeout:
            if self.channel.recv_ready():
                chunk = self.channel.recv(65535).decode('utf-8','ignore')
                output += chunk
            else:
                time.sleep(0.1)
        return output

    def _enter_system_view(self):
        self._send_command("sys")
        return "Error:Failed to enter system view" not in self._receive_output()

    def configuration_remote_management(self):
        if not self._enter_system_view():
            return False

        commands = [
            "aaa",
            f"{self.command_map['password policy']} administrator",
            "password history record number 0",
            "quit",
            f"local-user {self.username} password irreversible-cipher {self.passwd}",
            f"local-user {self.username} privilege level 15",
            f"local-user {self.username} service-type ssh telnet",
            "quit",
            "stelnet server enable",
            "user-interface vty 0 4",
            "authentication-mode aaa",
            "protocol in bound all",
            "quit",
            self.command_map['rsa_create'],
            "y"
        ]